{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f2c3496f-5fc3-4975-a812-f5a066dda84f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import glob\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "import fiftyone as fo\n",
    "\n",
    "from tqdm import tqdm\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools import mask as maskUtils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c5e0de57-f91f-41a0-bce0-23d46105802b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name=\"nike\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9cb48310-c131-4670-8710-bd909e39de74",
   "metadata": {},
   "outputs": [],
   "source": [
    "fo.delete_datasets('*')\n",
    "\n",
    "dataset = fo.Dataset(name=dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1769f529-4f05-43b8-9ddc-37edc4367ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = '/data/pt/data/NikeDatasets/increasing/xiantou'\n",
    "img_root_dir = osp.join(root_dir, \"images\")\n",
    "anno_dir = osp.join(root_dir, \"annotations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0ea57974-a05f-4740-88d4-e92c7b5761ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def cal_coordinate(bbox, offset, ori_shape):\n",
    "    x, y, w, h = list(map(int, bbox))\n",
    "    ori_width= ori_shape[1]\n",
    "    ori_height =ori_shape[0]\n",
    "\n",
    "    x1 = max(0, x-offset*ori_width)\n",
    "    x2 = min(x+w+offset*ori_width, ori_width)\n",
    "    y1 = max(0, y-offset*ori_height)\n",
    "    y2 = min(y+h+offset*ori_height, ori_height)\n",
    "    return list(map(int, (x1, y1, x2, y2)))\n",
    "\n",
    "def rename(path, pref=\"F\"):\n",
    "    file_name = osp.basename(path)\n",
    "    parent = osp.dirname(path)\n",
    "    file_name = \"_\".join([pref, file_name])\n",
    "    print(file_name, parent)\n",
    "    return osp.join(parent, file_name)\n",
    "\n",
    "def get_cat_id(cat_name):\n",
    "    if cat_name in [\"线头\",\"xiantou\"]:\n",
    "        return 1\n",
    "    if cat_name in [\"脏污\",\"zangwu\"]:\n",
    "        return 2\n",
    "    if cat_name in [\"溢胶\",\"yijiao\"]:\n",
    "        return 0\n",
    "\n",
    "def save_yolo_annotations(yolo_data, save_directory):\n",
    "    if not os.path.exists(save_directory):\n",
    "        os.makedirs(save_directory)\n",
    "    for file_name, annotations in yolo_data.items():\n",
    "        with open(os.path.join(save_directory, file_name.replace('.png', '.txt')), \"w\") as f:\n",
    "            for annotation in annotations:\n",
    "                f.write(annotation + \"\\n\")\n",
    "\n",
    "def get_annotations(annotations, categories, img_shape):\n",
    "    category_dict = {cat[\"id\"]: cat[\"name\"] for cat in categories}\n",
    "    polylines = []\n",
    "    detections = []\n",
    "    img_h, img_w = img_shape\n",
    "    for ann in annotations:\n",
    "        category_id = int(ann['category_id']) \n",
    "        coordinates = np.array(ann['segmentation'])\n",
    "        bbox = list(map(float, ann['bbox']))\n",
    "        x,y,w,h = [bbox[0]/img_w, bbox[1]/img_h, bbox[2]/img_w, bbox[3]/img_h]\n",
    "\n",
    "        polygon_points = coordinates / np.array([img_w, img_h])\n",
    "        category_name = category_dict[category_id]\n",
    "        polylines.append(fo.Polyline(\n",
    "            points=[polygon_points],\n",
    "            closed=True,\n",
    "            filled=True,\n",
    "        ))       \n",
    "        detections.append(fo.Detection(label=category_name, bounding_box=[x, y, w, h]))\n",
    "    return polylines, detections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "19b28a4a-137f-4e15-9d42-a4804bff542f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.02s)\n",
      "creating index...\n",
      "index created!\n",
      "all categories:  xiantou\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(10, 2)\n",
      "(10, 2)\n",
      "(18, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(16, 2)\n",
      "(10, 2)\n",
      "(11, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(12, 2)\n",
      "(7, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(6, 2)\n",
      "(8, 2)\n",
      "(12, 2)\n",
      "(11, 2)\n",
      "(7, 2)\n",
      "(12, 2)\n",
      "(10, 2)\n",
      "(14, 2)\n",
      "(11, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(8, 2)\n",
      "(9, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(10, 2)\n",
      "(13, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(13, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(10, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(10, 2)\n",
      "(14, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(13, 2)\n",
      "(20, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(21, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(6, 2)\n",
      "(10, 2)\n",
      "(15, 2)\n",
      "(15, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(13, 2)\n",
      "(20, 2)\n",
      "(10, 2)\n",
      "(13, 2)\n",
      "(10, 2)\n",
      "(15, 2)\n",
      "(12, 2)\n",
      "(11, 2)\n",
      "(8, 2)\n",
      "(8, 2)\n",
      "(35, 2)\n",
      "(16, 2)\n",
      "(13, 2)\n",
      "(11, 2)\n",
      "(11, 2)\n",
      "(17, 2)\n",
      "(8, 2)\n",
      "(7, 2)\n",
      "(6, 2)\n",
      "(5, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(16, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(16, 2)\n",
      "(13, 2)\n",
      "(11, 2)\n",
      "(7, 2)\n",
      "(18, 2)\n",
      "(11, 2)\n",
      "(7, 2)\n",
      "(17, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(15, 2)\n",
      "(8, 2)\n",
      "(11, 2)\n",
      "(14, 2)\n",
      "(11, 2)\n",
      "(4, 2)\n",
      "(4, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(15, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(9, 2)\n",
      "(20, 2)\n",
      "(13, 2)\n",
      "(15, 2)\n",
      "(11, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(8, 2)\n",
      "(11, 2)\n",
      "(16, 2)\n",
      "(8, 2)\n",
      "(10, 2)\n",
      "(7, 2)\n",
      "(14, 2)\n",
      "(12, 2)\n",
      "(11, 2)\n",
      "(16, 2)\n",
      "(10, 2)\n",
      "(8, 2)\n",
      "(12, 2)\n",
      "(7, 2)\n",
      "(14, 2)\n",
      "(13, 2)\n",
      "(8, 2)\n",
      "(16, 2)\n",
      "(7, 2)\n",
      "(7, 2)\n",
      "(7, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(14, 2)\n",
      "(10, 2)\n",
      "(11, 2)\n",
      "(21, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(8, 2)\n",
      "(9, 2)\n",
      "(18, 2)\n",
      "(21, 2)\n",
      "(12, 2)\n",
      "(10, 2)\n",
      "(6, 2)\n",
      "(18, 2)\n",
      "(13, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(10, 2)\n",
      "(13, 2)\n",
      "(8, 2)\n",
      "(6, 2)\n",
      "(10, 2)\n",
      "(17, 2)\n",
      "(13, 2)\n",
      "(14, 2)\n",
      "(16, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(12, 2)\n",
      "(13, 2)\n",
      "(7, 2)\n",
      "(9, 2)\n",
      "(6, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(12, 2)\n",
      "(6, 2)\n",
      "(8, 2)\n",
      "(12, 2)\n",
      "(8, 2)\n",
      "(27, 2)\n",
      "(11, 2)\n",
      "(7, 2)\n",
      "(5, 2)\n",
      "(11, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(8, 2)\n",
      "(19, 2)\n",
      "(11, 2)\n",
      "(13, 2)\n",
      "(16, 2)\n",
      "(18, 2)\n",
      "(6, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(7, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(9, 2)\n",
      "(15, 2)\n",
      "(7, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(19, 2)\n",
      "(13, 2)\n",
      "(16, 2)\n",
      "(10, 2)\n",
      "(7, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(15, 2)\n",
      "(10, 2)\n",
      "(13, 2)\n",
      "(19, 2)\n",
      "(10, 2)\n",
      "(15, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(8, 2)\n",
      "(14, 2)\n",
      "(16, 2)\n",
      "(12, 2)\n",
      "(25, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(12, 2)\n",
      "(17, 2)\n",
      "(14, 2)\n",
      "(10, 2)\n",
      "(15, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(29, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(12, 2)\n",
      "(17, 2)\n",
      "(7, 2)\n",
      "(12, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(11, 2)\n",
      "(8, 2)\n",
      "(8, 2)\n",
      "(9, 2)\n",
      "(6, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(18, 2)\n",
      "(9, 2)\n",
      "(15, 2)\n",
      "(8, 2)\n",
      "(7, 2)\n",
      "(12, 2)\n",
      "(8, 2)\n",
      "(8, 2)\n",
      "(13, 2)\n",
      "(9, 2)\n",
      "(16, 2)\n",
      "(18, 2)\n",
      "(19, 2)\n",
      "(19, 2)\n",
      "(10, 2)\n",
      "(11, 2)\n",
      "(18, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(14, 2)\n",
      "(17, 2)\n",
      "(15, 2)\n",
      "(10, 2)\n",
      "(16, 2)\n",
      "(12, 2)\n",
      "(14, 2)\n",
      "(12, 2)\n",
      "(11, 2)\n",
      "(8, 2)\n",
      "(8, 2)\n",
      "(10, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(8, 2)\n",
      "(11, 2)\n",
      "(8, 2)\n",
      "(10, 2)\n",
      "(6, 2)\n",
      "(9, 2)\n",
      "(15, 2)\n",
      "(10, 2)\n",
      "(14, 2)\n",
      "(15, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(10, 2)\n",
      "(14, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(14, 2)\n",
      "(10, 2)\n",
      "(12, 2)\n",
      "(33, 2)\n",
      "(19, 2)\n",
      "(18, 2)\n",
      "(6, 2)\n",
      "(6, 2)\n",
      "(5, 2)\n",
      "(4, 2)\n",
      "(8, 2)\n",
      "(10, 2)\n",
      "(57, 2)\n",
      "(8, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(9, 2)\n",
      "(17, 2)\n",
      "(9, 2)\n",
      "(10, 2)\n",
      "(12, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(13, 2)\n",
      "(7, 2)\n",
      "(7, 2)\n",
      "(8, 2)\n",
      "(11, 2)\n",
      "(10, 2)\n",
      "(7, 2)\n",
      "(17, 2)\n",
      "(11, 2)\n",
      "(17, 2)\n",
      "(7, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(13, 2)\n",
      "(21, 2)\n",
      "(5, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(12, 2)\n",
      "(9, 2)\n",
      "(11, 2)\n",
      "(12, 2)\n",
      "(10, 2)\n",
      "(9, 2)\n",
      "(5, 2)\n"
     ]
    }
   ],
   "source": [
    "without_shoes = []\n",
    "\n",
    "duplicate_file = []\n",
    "img_not_found = []\n",
    "illegal_img = []\n",
    "save_annotations = True\n",
    "save_img = False\n",
    "is_crop = False\n",
    "\n",
    "defect_name = [\"线头\", \"溢胶\", \"脏污\"]\n",
    "\n",
    "defect_name = [\"xiantou\"]\n",
    "\n",
    "fo_samples = []\n",
    "\n",
    "for idx, json_file in enumerate(glob.glob(f\"{anno_dir}/**/*.json\", recursive=True)) :\n",
    "    coco = COCO(json_file)\n",
    "    categories = coco.loadCats(coco.getCatIds())\n",
    "    print(\"all categories: \", \", \".join([cat[\"name\"] for cat in categories ]))\n",
    "\n",
    "    for img in coco.imgs.values():\n",
    "        img_path = glob.glob(f\"{img_root_dir}/**/{img['file_name']}\", recursive=True)\n",
    "        if len(img_path) > 1:\n",
    "            duplicate_file.append(json_file)\n",
    "            break\n",
    "        if len(img_path) == 0:\n",
    "            img_not_found.append(img_path)\n",
    "            continue\n",
    "        if img['height'] == 0 or img['width'] == 0:\n",
    "            illegal_img.append(img_path)\n",
    "            continue\n",
    "        img_path = img_path[0]\n",
    "\n",
    "        real_path = os.path.realpath(img_path)\n",
    "\n",
    "        fo_sample = fo.Sample(filepath=real_path) \n",
    "        polylines = []\n",
    "        detections = []\n",
    "        \n",
    "        img_w = img['width']\n",
    "        img_h = img['height']\n",
    "    \n",
    "        cat_id = [cat[\"id\"] for cat in categories if cat['name'] in defect_name]\n",
    "        annotations = coco.loadAnns(coco.getAnnIds(imgIds=img[\"id\"], catIds=cat_id, iscrowd=None))\n",
    "        polylines, detections = get_annotations(annotations, coco.dataset[\"categories\"], (img_h, img_w))\n",
    "        fo_sample[\"prediction\"] = fo.Detections(detections=detections)\n",
    "        fo_sample[\"polylines\"] = fo.Polylines(polylines=polylines)\n",
    "        fo_samples.append(fo_sample)      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "f9cc56cd-792a-4e18-906c-3a1a7fbe3064",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "print(len(img_not_found))\n",
    "print(len(illegal_img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "aab352da-6176-4d09-b41b-edb12c3ef754",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 100% |███████████████| 1869/1869 [1.1s elapsed, 0s remaining, 1.7K samples/s]         \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"800\"\n",
       "            src=\"http://0.0.0.0:8989/?notebook=True&subscription=a290fa32-4f7e-44bb-9345-ca866249ac8f\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7f34c8690290>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset.add_samples(fo_samples)\n",
    "session = fo.launch_app(dataset, address=\"0.0.0.0\", port=8989)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e299ff40-41bf-4708-af8d-67eedca3086e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
